package installer;

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;

import testing.Notifier;
import testing.PredictionTester;
import data.DataInfo;

/**
 * Loads the data from the original download folder, Generates extra information
 * (userIDs, etc), converts to user and movie indexed slabs, and saves it all in
 * easily loaded files.
 * 
 * @author Robert Layton
 * 
 */
public class Installer {

	private HashMap<Integer, Integer> userIndices = new HashMap<Integer, Integer>();

	private String downloadFolder;

	private String outputFolder;

	private boolean getDate;

	private boolean movieIndexedLoaded = false;

	private boolean userIndexedLoaded = false;

	private int[][] movieIndexedUserIDs = new int[DataInfo.NUM_MOVIES][];

	private byte[][] movieIndexedRatings = new byte[DataInfo.NUM_MOVIES][];

	private byte[][][] movieIndexedDates = new byte[DataInfo.NUM_MOVIES][][];

	private short[][] userIndexedMovieIDs;

	private byte[][] userIndexedRatings;

	private byte[][][] userIndexedDates;

	private int[] userIDs;

	// Probe data
	private boolean probeScrubbed = false;

	private static int[] probeUserIDs;

	private static short[] probeMovieIDs;

	private static byte[] probeRatings;

	public Installer(String downloadFolder, String outputFolder, boolean getDate) {
		this.downloadFolder = downloadFolder;
		this.outputFolder = outputFolder;
		this.getDate = getDate;
	}

	public void install() throws IOException {
		System.out.println("Beginning Data Check and Installation procedure");
		createMovieIndexedRatings();
		createUserIndexedRatings();
		createAveragesFiles();
		testData();

		System.out.println("With probe dataset complete.");
		scrubProbe();

		System.out.println("All data files should now be installed.");
	}

	private void scrubProbe() {
		// Remove the probe data
		loadTrainingData();
		loadUserIndexedRatings();
		loadUserIndices();

		final int trainingLength = probeUserIDs.length;
		Notifier n = new Notifier(trainingLength,
				"Scrubbing probe from data: [", "]");
		for (int i = 0; i < trainingLength; i++) {
			short movieID = probeMovieIDs[i];
			int userID = probeUserIDs[i];
			removeRating(movieID, userID);
			n.next();
		}

		probeScrubbed = true;
		blitMovies();
		blitUsers();
	}

	private void removeRating(short movieID, int userID) {
		// Remove by movie-indexed
		int userIndex = userIndices.get(userID);
		int[] userIDs = movieIndexedUserIDs[movieID - 1];
		
		byte[] mRatings = movieIndexedRatings[movieID - 1];
		int userI = Arrays.binarySearch(userIDs, userIndex);

		movieIndexedUserIDs[movieID - 1] = new int[userIDs.length - 1];
		movieIndexedRatings[movieID - 1] = new byte[userIDs.length - 1];

		int curIndex = 0;
		for (int i = 0; i < userIDs.length; i++) {
			if (i != userI) {
				movieIndexedUserIDs[movieID - 1][curIndex] = userIDs[i];
				movieIndexedRatings[movieID - 1][curIndex] = mRatings[i];
				curIndex++;
			}
		}
		if (curIndex != movieIndexedUserIDs[movieID - 1].length) {
			System.err.println("Invalid number of entries (movie): " + curIndex
					+ "/" + movieIndexedUserIDs[movieID - 1].length);
		}
		// Remove by user-indexed
		short[] movieIDs = userIndexedMovieIDs[userIndex];
		byte[] uRatings = userIndexedRatings[userIndex];
		int movieI = Arrays.binarySearch(movieIDs, movieID);

		userIndexedMovieIDs[userIndex] = new short[movieIDs.length - 1];
		userIndexedRatings[userIndex] = new byte[movieIDs.length - 1];

		curIndex = 0;
		for (int i = 0; i < movieIDs.length; i++) {
			if (i != movieI) {
				userIndexedMovieIDs[userIndex][curIndex] = movieIDs[i];
				userIndexedRatings[userIndex][curIndex] = uRatings[i];
				curIndex++;
			}
		}
		if (curIndex != userIndexedMovieIDs[userIndex].length) {
			System.err.println("Invalid number of entries (user): " + curIndex
					+ "/" + userIndexedMovieIDs[userIndex].length);
		}
	}

	private void loadTrainingData() {
		String filename = DataInfo.downloadFolder + "betterprobe.txt";
		File probeFile = new File(filename);
		// 1425333 is the number of ratings in the probe file
		int numProbe = 1408395;
		Notifier n = new Notifier(numProbe, "Loading probe set: [", "]");
		int count = 0;

		probeUserIDs = new int[numProbe];

		probeMovieIDs = new short[numProbe];

		probeRatings = new byte[numProbe];
		try {
			BufferedReader in = new BufferedReader(new FileReader(probeFile));
			if (!in.ready())
				throw new IOException();
			// First line is the movie number: "200:\n"

			String line = "";
			// Get the movieID only
			short movieID = 0;
			int userID = 0;
			// System.out.print("Movie: " + movieID);
			while ((line = in.readLine()) != null) {
				// If the line is a new movieID
				if (line.contains(":")) {
					movieID = Short.parseShort(line.substring(0,
							line.length() - 1));
				} else {
					// Create the ratings
					ArrayList<String> tokens = PredictionTester.splitLine(line,
							',');
					userID = Integer.parseInt(tokens.get(1));
					byte rating = Byte.parseByte(tokens.get(0));
					if (userID < 1) {
						System.err.println("Invalid user ID: " + userID);
						System.err.println("Line: " + line);
					}
					// Store the rating in memory
					probeUserIDs[count] = userID;

					probeMovieIDs[count] = movieID;

					probeRatings[count] = rating;
					n.next();
					count++;
				}

			}

			in.close();
		} catch (IOException e) {
			System.out.println(e);
		}
		System.out.println("Loaded " + count + " entries in the probe set.");
	}

	private void createMovieIndexedRatings() {
		// Check if the movie Indexed slab has been created
		boolean movieFileCreated = false;

		String filename = outputFolder + "movieIndexedSlab.data";
		File f2 = new File(filename);
		movieFileCreated = f2.exists();
		if (!movieFileCreated) {
			System.out
					.println("Movie indexed file does not exist. Creating it now.");
			// If the file hasn't been created, then create it
			// Load all files into movie-indexed arrays
			File folder = new File(downloadFolder + "/training_set");
			// For each file in this folder
			File[] listOfFiles = folder.listFiles();
			System.out.println("There are " + listOfFiles.length
					+ " files to load.");
			Notifier n = new Notifier(DataInfo.NUM_MOVIES, "Loading by movie "
					+ ": [", "]");
			for (File f : listOfFiles) {
				// parse each file
				parseFile(f, getDate);
				n.next();
			}

			createUserIndices();
			blitMovies();
		} else {
			System.out.println("Movie indexed ratings already created.");
			loadMovieIndexedRatings();
		}
		movieIndexedLoaded = true;
	}

	private void createUserIndices() {
		Notifier n = new Notifier(DataInfo.NUM_MOVIES,
				"Obtaining user indexes [", "]");

		// Iterate through all the ratings we have, getting new indices as we go
		HashSet<Integer> allUserIDs = new HashSet<Integer>();
		for (short movieID = 1; movieID <= DataInfo.NUM_MOVIES; movieID++) {
			int L = movieIndexedUserIDs[movieID - 1].length;
			for (int j = 0; j < L; j++) {
				Integer userID = movieIndexedUserIDs[movieID - 1][j];
				allUserIDs.add(userID);
			}
			n.next();
		}
		// Check we have the correct number of user IDs
		if (allUserIDs.size() != DataInfo.NUM_USERS) {
			int difference = DataInfo.NUM_USERS - allUserIDs.size();
			System.err.println("Invalid number of user IDs found: delta of "
					+ difference);
		}

		// Find all the unique userIDs
		userIDs = new int[allUserIDs.size()];
		int index = 0;
		for (int userID : allUserIDs) {
			userIDs[index] = userID;
			index++;
		}

		// Sort the user indexes
		Arrays.sort(userIDs);

		// Create a new hashmap for the sequential IDs
		userIndices = new HashMap<Integer, Integer>();
		for (int i = 0; i < userIDs.length; i++) {
			// this gives each value a sequential ID
			userIndices.put(userIDs[i], i);
		}

		// Save the user IDs and their indexes to a file
		String indexfilename = outputFolder + "userIndexes.data";
		System.out.print("Blitting indexes to file...");
		BufferedOutputStream out = null;
		ObjectOutputStream oos = null;
		try {
			out = new BufferedOutputStream(new FileOutputStream(indexfilename));
			oos = new ObjectOutputStream(out);
			// Write the arrays to the file
			oos.writeObject(userIDs);
			oos.close();
			out.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		System.out.println("done");

		// Iterate through all ratings, replacing userID with userIndex
		n = new Notifier(DataInfo.NUM_MOVIES,
				"Replacing user IDs with indices [", "]");
		for (short movieID = 0; movieID < DataInfo.NUM_MOVIES; movieID++) {
			int L = movieIndexedUserIDs[movieID].length;
			for (int i = 0; i < L; i++) {
				int ui = userIndices.get(movieIndexedUserIDs[movieID][i]);
				if (ui == -1) {
					throw new RuntimeException("Invalid user index: "
							+ movieIndexedUserIDs[movieID][i]);
				}
				movieIndexedUserIDs[movieID][i] = ui;
			}
			n.next();
		}
	}

	private void createUserIndexedRatings() {

		String filename = outputFolder + "userIndexedSlab.data";
		if (getDate) {
			filename = outputFolder + "userIndexedSlabDates.data";
		}
		File f2 = new File(filename);
		boolean userFileCreated = f2.exists();
		if (userFileCreated) {
			System.out.println("User indexed file has already been created");
			loadUserIndexedRatings();
			return;
		} else {
			System.out
					.println("User indexed file not created, creating it now");
			// Load movie indexed file to memory
			loadMovieIndexedRatings();

			// Get the user indexes here from file
			loadUserIndices();

			// Find out how many ratings each user has
			System.out.print("Obtaining rating counts...");
			int[] numRatings = new int[DataInfo.NUM_USERS];
			for (int i = 0; i < numRatings.length; i++) {
				numRatings[i] = 0;
			}
			for (short movieID = 1; movieID <= DataInfo.NUM_MOVIES; movieID++) {
				for (int i = 0; i < movieIndexedUserIDs[movieID - 1].length; i++) {
					int userIndex = movieIndexedUserIDs[movieID - 1][i];
					// System.out.println(userIndex);
					numRatings[userIndex]++;
				}
			}
			System.out.println("done");

			// Create the arrays that the ratings will go in.
			System.out.print("Initializing arrays...");
			userIndexedMovieIDs = new short[DataInfo.NUM_USERS][];
			userIndexedRatings = new byte[DataInfo.NUM_USERS][];
			for (int i = 0; i < userIDs.length; i++) {
				userIndexedMovieIDs[i] = new short[numRatings[i]];
				userIndexedRatings[i] = new byte[numRatings[i]];
			}

			// Index stores where we are up to with each uesr's ratings
			int[] index = new int[DataInfo.NUM_USERS];
			for (int i = 0; i < index.length; i++) {
				index[i] = 0;
			}
			System.out.println("done");

			// Now put all of the ratings into the proper user-indexed buckets
			System.out.print("Indexing by user...");
			for (short movieID = 1; movieID <= DataInfo.NUM_MOVIES; movieID++) {
				for (int i = 0; i < movieIndexedUserIDs[movieID - 1].length; i++) {
					int userIndex = movieIndexedUserIDs[movieID - 1][i];
					byte rating = movieIndexedRatings[movieID - 1][i];
					// Save the rating to user indexed
					userIndexedMovieIDs[userIndex][index[userIndex]] = (short) (movieID);
					userIndexedRatings[userIndex][index[userIndex]] = rating;
					index[userIndex]++;
				}
			}
			System.out.println("done");

			// Verify that we found all ratings
			System.out.print("Verifying indexes...");
			for (int i = 0; i < numRatings.length; i++) {
				if (index[i] != numRatings[i] || numRatings[i] == 0) {
					System.err
							.println("Did not get enough ratings for user index: "
									+ i);
				}
			}
			System.out.println("done.");

			System.out.print("Sorting indexes...");
			for (short movieIndex = 0; movieIndex < DataInfo.NUM_MOVIES; movieIndex++) {
				// Sort the user indexes for each movie's set of raters
				int[] raters = movieIndexedUserIDs[movieIndex];
				byte[] ratings = movieIndexedRatings[movieIndex];
				HashMap<Integer, Byte> mRatings = new HashMap<Integer, Byte>();
				for (int i = 0; i < raters.length; i++) {
					mRatings.put(raters[i], ratings[i]);
				}
				Arrays.sort(raters);
				movieIndexedUserIDs[movieIndex] = raters;
				for (int i = 0; i < raters.length; i++) {
					movieIndexedRatings[movieIndex][i] = mRatings
							.get(raters[i]);
				}
			}
			for (int userIndex = 0; userIndex < DataInfo.NUM_USERS; userIndex++) {
				// sort the movie ids for each user's rated movies
				short[] ratedMovies = userIndexedMovieIDs[userIndex];
				byte[] ratings = userIndexedRatings[userIndex];
				HashMap<Short, Byte> uRatings = new HashMap<Short, Byte>();
				for (int i = 0; i < ratedMovies.length; i++) {
					uRatings.put(ratedMovies[i], ratings[i]);
				}
				Arrays.sort(ratedMovies);
				userIndexedMovieIDs[userIndex] = ratedMovies;
				for (int i = 0; i < ratedMovies.length; i++) {
					userIndexedRatings[userIndex][i] = uRatings
							.get(ratedMovies[i]);
				}
			}
			System.out.println("done");

			// Blit the user indexed to file
			blitUsers();
			// Blit the movie indexed file again (it has been sorted now)
			blitMovies();
		}
	}

	private void createAveragesFiles() {
		String filename = outputFolder + "averagesSlab.data";
		File f2 = new File(filename);
		if (f2.exists()) {
			System.out.println("Averages file exists");
			return;
		}

		System.out.println("Averages file does not exist, creating it now...");
		loadMovieIndexedRatings();
		loadUserIndexedRatings();
		System.out.print("Calculating averages...");

		// user info
		double[] userAverages = new double[DataInfo.NUM_USERS];
		double[] userStdDev = new double[DataInfo.NUM_USERS];
		for (int u = 0; u < DataInfo.NUM_USERS; u++) {
			byte[] ratings = userIndexedRatings[u];
			userAverages[u] = average(ratings);
			userStdDev[u] = stddev(ratings, userAverages[u]);
			DataInfo.checkInvalidDouble("userAverages[" + u + "]",
					userAverages[u]);
			DataInfo.checkInvalidDouble("userStdDev[" + u + "]", userStdDev[u]);
		}

		// Movie info
		double[] movieAverages = new double[DataInfo.NUM_MOVIES];
		double[] movieStdDev = new double[DataInfo.NUM_MOVIES];
		for (int m = 0; m < DataInfo.NUM_MOVIES; m++) {
			byte[] ratings = movieIndexedRatings[m];
			movieAverages[m] = average(ratings);
			movieStdDev[m] = stddev(ratings, movieAverages[m]);
			DataInfo.checkInvalidDouble("movieAverages[" + m + "]",
					movieAverages[m]);
			DataInfo.checkInvalidDouble("movieStdDev[" + m + "]",
					movieStdDev[m]);
		}
		System.out.println("done");
		System.out.print("Blitting averages...");
		BufferedOutputStream out = null;
		ObjectOutputStream oos = null;
		try {
			// Filename from top of method
			out = new BufferedOutputStream(new FileOutputStream(filename));
			oos = new ObjectOutputStream(out);
			// Write the arrays to the file
			oos.writeObject(userAverages);
			oos.writeObject(movieAverages);
			oos.writeObject(userStdDev);
			oos.writeObject(movieStdDev);
			oos.close();
			out.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		System.out.println("done.");
	}

	private double average(byte[] ratings) {
		if (ratings.length == 0)
			return 2.5; // middle
		double sum = 0.0;
		for (byte r : ratings)
			sum += r;
		return sum / ratings.length;
	}

	private double stddev(byte[] ratings, double mean) {
		if (ratings.length == 0)
			return 0;
		double sum = 0.0;
		for (byte r : ratings) {
			double e = mean - r;
			sum += e * e;
		}
		double stddev = Math.sqrt(sum / (double) ratings.length);
		return stddev;
	}

	private void loadUserIndexedRatings() {
		if (userIndexedLoaded) {
			return;
		}
		ObjectInputStream ois = null;
		System.out.print("Loading user indexed ratings...");
		try {
			String movieFilename = DataInfo.binaryFolder
					+ "userIndexedSlab.data";
			ois = new ObjectInputStream(new FileInputStream(movieFilename));
			userIndexedMovieIDs = (short[][]) ois.readObject();
			userIndexedRatings = (byte[][]) ois.readObject();
			ois.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		}
		System.out.println("done");
		userIndexedLoaded = true;
	}

	private void blitMovies() {
		String filename = outputFolder + "movieIndexedSlab.data";
		if (probeScrubbed && !getDate) {
			filename = outputFolder + "cleanedMovieIndexedSlab.data";
		} else if (getDate) {
			if (probeScrubbed) {
				filename = outputFolder + "cleanedMovieIndexedSlabDates.data";
			} else {
				filename = outputFolder + "movieIndexedSlabDates.data";
			}
		}
		System.out.print("Blitting to movie index to file...");
		BufferedOutputStream out = null;
		ObjectOutputStream oos = null;
		try {
			out = new BufferedOutputStream(new FileOutputStream(filename));
			oos = new ObjectOutputStream(out);
			// Write the arrays to the file
			oos.writeObject(movieIndexedUserIDs);
			oos.writeObject(movieIndexedRatings);
			if (getDate) {
				oos.writeObject(movieIndexedDates);
			}
			oos.close();
			out.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		System.out.println("done");
	}

	private void parseFile(File f, boolean getDate) {
		try {
			BufferedReader in = new BufferedReader(new FileReader(f));
			if (!in.ready())
				throw new IOException();
			// First line is the movie number: "200:\n"
			String line = in.readLine();
			// Get the movieID only
			short movieID = Short.parseShort(line.substring(0,
					line.length() - 1));

			// Load all of the rest of the lines into memory
			ArrayList<String> lines = new ArrayList<String>();
			// System.out.print("Movie: " + movieID);
			while ((line = in.readLine()) != null) {
				lines.add(line);
			}

			// Setup the array
			movieIndexedUserIDs[movieID - 1] = new int[lines.size()];
			movieIndexedRatings[movieID - 1] = new byte[lines.size()];
			if (getDate) {
				movieIndexedDates[movieID - 1] = new byte[lines.size()][3];
			}

			int i = 0;
			for (String curLine : lines) {
				// Create the ratings
				ArrayList<String> tokens = splitLine(curLine, ',');
				movieIndexedUserIDs[movieID - 1][i] = Integer.parseInt(tokens
						.get(0));
				movieIndexedRatings[movieID - 1][i] = Byte.parseByte(tokens
						.get(1));
				if (getDate) {
					ArrayList<String> dateTokens = splitLine(tokens.get(2), '-');
					movieIndexedDates[movieID - 1][i][0] = Byte
							.parseByte(dateTokens.get(0).substring(2));
					movieIndexedDates[movieID - 1][i][1] = Byte
							.parseByte(dateTokens.get(1));
					movieIndexedDates[movieID - 1][i][2] = Byte
							.parseByte(dateTokens.get(2));
				}
				i++;
			}

			in.close();
		} catch (IOException e) {
			System.out.println(e);
		}
	}

	private ArrayList<String> splitLine(String line, char separator) {
		ArrayList<String> tokens = new ArrayList<String>();
		String cur = "";
		for (char c : line.toCharArray()) {
			if (c == separator) {
				// Separator found, new token
				tokens.add(cur);
				cur = "";
			} else {
				// Append to cur String
				cur += c;
			}
		}
		// Add the last token
		tokens.add(cur);
		return tokens;
	}

	private void loadMovieIndexedRatings() {
		if (movieIndexedLoaded) {
			return;
		}
		ObjectInputStream ois = null;
		System.out.print("Loading movie indexed ratings...");
		try {
			String movieFilename = DataInfo.binaryFolder
					+ "movieIndexedSlab.data";
			ois = new ObjectInputStream(new FileInputStream(movieFilename));
			movieIndexedUserIDs = (int[][]) ois.readObject();
			movieIndexedRatings = (byte[][]) ois.readObject();
			ois.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		}
		System.out.println("done");
		movieIndexedLoaded = true;

		// Find out if the userIDs have been swapped with user indexes or not
		boolean highFound = false;
		for (int m = 0; m < DataInfo.NUM_MOVIES; m++) {
			int size = movieIndexedUserIDs[m].length;
			for (int u = 0; u < size; u++) {
				if (movieIndexedUserIDs[m][u] > DataInfo.NUM_USERS) {
					highFound = true;
				}
			}
		}
		if (highFound) {
			System.err
					.println("High user ID found, haven't been swapped with user indexes yet");
		} else {
			System.out.println("User indices have been replaced in the data");
		}
	}

	private void loadUserIndices() {
		ObjectInputStream ois = null;
		System.out.print("Loading movie indexed ratings...");
		try {
			String userIndexesFilename = DataInfo.binaryFolder
					+ "userIndexes.data";
			ois = new ObjectInputStream(
					new FileInputStream(userIndexesFilename));
			this.userIDs = (int[]) ois.readObject();
			ois.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		}
		System.out.println("done");

		// Now load them into our hashMap for fast retrieval
		userIndices = new HashMap<Integer, Integer>();
		int i = 0;
		for (int userID : userIDs) {
			if (userID < 0) {
				throw new RuntimeException(
						"Invalid User ID found in index file");
			}
			userIndices.put(userID, i++);
		}
	}

	private void blitUsers() {
		String filename = outputFolder + "userIndexedSlab.data";
		if (probeScrubbed && !getDate) {
			filename = outputFolder + "cleanedUserIndexedSlab.data";
		} else if (getDate) {
			if (!probeScrubbed) {
				filename = outputFolder + "userIndexedSlabDates.data";
			} else {
				filename = outputFolder + "cleanedUserIndexedSlabDates.data";
			}
		}
		System.out.print("Blitting to user index file...");
		BufferedOutputStream out = null;
		ObjectOutputStream oos = null;
		try {
			out = new BufferedOutputStream(new FileOutputStream(filename));
			oos = new ObjectOutputStream(out);
			// Write the arrays to the file
			oos.writeObject(userIndexedMovieIDs);
			userIndexedMovieIDs = null;
			System.gc();
			oos.writeObject(userIndexedRatings);
			userIndexedRatings = null;
			System.gc();
			if (getDate) {
				oos.writeObject(userIndexedDates);
			}
			oos.close();
			out.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		System.out.println("done");
	}

	private boolean testData() {
		return true;
	}

	public static void main(String args[]) {
		boolean getDate = false;
		String downloadFolder = "/home/bob/projects/netflix/data/";
		String outputFolder = "/home/bob/projects/netflix/binary/";

		Installer i = new Installer(downloadFolder, outputFolder, getDate);
		try {
			i.install();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
}
